#!/usr/bin/env python3

import subprocess
import os
import json
from reportlab.lib.pagesizes import letter
from reportlab.platypus import SimpleDocTemplate, Paragraph, Image
from reportlab.lib.styles import getSampleStyleSheet
from datetime import datetime
import json

import matplotlib.pyplot as plt
from DashscopeAPI import DASHSCOPE_IMPORT_SUCCESS

from pylatex import Document, Section, Subsection, Command, Figure, Subsubsection
from pylatex.utils import italic, NoEscape
import re


def ReadJson(json_file):
    with open(json_file, "r") as f:
        data = json.load(f)
    return data


def ExtractDateFromFilename(filename):
    pattern = re.compile(r"[\-_](\d{4}-\d{2}-\d{2}-\d{2}-\d{2}-\d{2})")
    match = pattern.search(filename)
    if not match:
        return None
    date_string = match.group(1)
    # parts = os.path.splitext(filename)[0].split("_")
    # date_string = parts[-1]
    try:
        date_object = datetime.strptime(date_string, "%Y-%m-%d-%H-%M-%S")
        datetime_str = datetime.strftime(date_object, "%Y-%m-%d-%H")
        return datetime_str
    except ValueError:
        return None


def ConvertOldVersion(data):
    old_data = {"result": data.get("result", ""), "comment": data.get("comment", "")}
    bag_path = data.get("bag", "")
    data = {"bag": bag_path, "RUN000": old_data}
    return data


def GenerateLatexPDF(date_data: dict, output_file):
    doc = Document()
    doc.preamble.append(Command("title", "Results"))
    doc.append(Section("Results"))
    date_no = 1
    success_count = 0
    fail_count = 0
    failure_comments = []
    ws_path = os.path.join(os.path.dirname(__file__), "latexws")
    if not os.path.exists(ws_path):
        print("latexws not found")
        return

    for date, data in date_data.items():
        if data == None or len(data) == 0:
            continue
        doc.append(Subsection(f"{date}"))
        data_no = 1
        for d in data:
            if "result" in d:
                d = ConvertOldVersion(d)
            run_index = 0
            bag_file = d.get("bag", "")
            for key, value in d.items():
                if key.startswith("RUN"):
                    run_index += 1
                else:
                    continue
                comment = ""
                if "comment" in value:
                    comment = value.get("comment", "")
                elif "description" in value:
                    comment = value.get("description", "")
                result = value.get("result", "")
                parmeters = value.get("parameters", "")
                if result == "Success" or result == "success":
                    success_count += 1
                elif result == "Failure" or result == "failure":
                    fail_count += 1
                    failure_comments.append(comment)
                bag_file = os.path.basename(bag_file)
                doc.append(Subsubsection(f"{bag_file}_RUN{run_index}"))
                doc.append(NoEscape(r"\textit{Result:} ") + italic(result))
                doc.append("\n")
                doc.append(NoEscape(r"\textit{comment:} ") + italic(comment))
                doc.append("\n")
                doc.append(NoEscape(r"\textit{Paramters:} "))
                doc.append("\n")
                if isinstance(parmeters, dict):
                    for key, value in parmeters.items():
                        parameter_str = f"{key}: {value}\n"
                        doc.append(NoEscape(italic(parameter_str)))
                data_no += 1
        date_no += 1
    total_count = success_count + fail_count

    if total_count > 0:
        doc.append(Section("Summarize"))
        doc.append(NoEscape(r"\textit{Total: } ") + italic(total_count))
        doc.append("\n")
        doc.append(NoEscape(r"\textit{Success:} ") + italic(success_count))
        doc.append("\n")
        doc.append(NoEscape(r"\textit{Failure:} ") + italic(fail_count))
        doc.append("\n")
        labels = ["Success", "Failure"]
        sizes = [success_count / total_count, fail_count / total_count]
        plt.figure(figsize=(6, 6))
        plt.pie(sizes, labels=labels, autopct="%1.1f%%", startangle=140)
        plt.axis("equal")
        plt.title("Success Rate")
        temp_fig_path = os.path.join(ws_path, "temp.png")
        plt.savefig(temp_fig_path)
        plt.close()
        with doc.create(Figure(position="h!")) as plot:
            plot.add_image("temp.png", width="400px")
            plot.add_caption("Success Rate")

    doc.generate_pdf(os.path.join(ws_path, "temp"), clean_tex=True, silent=True)
    os.remove(temp_fig_path)
    mv_cmd = f"mv {ws_path}/temp.pdf {output_file}"
    subprocess.run(mv_cmd, shell=True)


def GeneratePDF(date_data: dict, output_file):
    doc = SimpleDocTemplate(output_file, pagesize=letter)
    styles = getSampleStyleSheet()
    failure_comments = []
    content = []
    content.append(Paragraph("Results", styles["Title"]))
    date_no = 1
    success_count = 0
    fail_count = 0
    for date, data in date_data.items():
        if data == None or len(data) == 0:
            continue
        content.append(Paragraph(f"{date_no}. {date}", styles["Heading1"]))
        data_no = 1
        for d in data:
            bag_file = d.get("bag", "")
            comment = d.get("comment", "")
            result = d.get("result", "")
            if result == "Success" or result == "success":
                success_count += 1
            elif result == "Failure" or result == "failure":
                fail_count += 1
                if not comment == None and not comment == "":
                    failure_comments.append(comment)
            bag_file = os.path.basename(bag_file)
            content.append(Paragraph(f"{data_no}. {bag_file}", styles["Heading2"]))
            content.append(Paragraph(f"Result: {result}", styles["Normal"]))
            content.append(Paragraph(f"comment: {comment}", styles["Normal"]))
            # Add empty line between entries
            content.append(Paragraph("", styles["Normal"]))
            data_no += 1
        date_no += 1
    total_count = success_count + fail_count
    if total_count > 0:
        content.append(Paragraph(f"Summarize", styles["Heading1"]))
        content.append(
            Paragraph(
                f"Total: {total_count}, Success: {success_count}, Failure: {fail_count}",
                styles["Normal"],
            )
        )
        labels = ["Success", "Failure"]
        sizes = [success_count / total_count, fail_count / total_count]
        plt.figure(figsize=(6, 6))
        plt.pie(sizes, labels=labels, autopct="%1.1f%%", startangle=140)
        plt.axis("equal")
        plt.title("Success Rate")
        plt.savefig("temp.png")
        plt.close()
        img = Image("temp.png")
        img.drawHeight = 400
        img.drawWidth = 400
        content.append(img)
    doc.build(content)
    if os.path.exists("temp.png"):
        os.remove("temp.png")


def GetAllJson(input_dir):
    json_files = [
        os.path.join(input_dir, file)
        for file in os.listdir(input_dir)
        if file.endswith(".json")
    ]
    data = []
    for file in json_files:
        data.append(ReadJson(file))
    sorted(data, key=lambda x: x.get("bag"))
    date_data = {}
    for d in data:
        date = ExtractDateFromFilename(d.get("bag", ""))
        if data == None:
            continue
        if date in date_data:
            date_data[date].append(d)
        else:
            date_data[date] = [d]
    return date_data


def main(input_dir, output_file):
    date_data = GetAllJson(input_dir)
    GeneratePDF(date_data, output_file)


if __name__ == "__main__":
    config_file = os.path.join(os.path.dirname(__file__), "config/default.json")
    config = ReadJson(config_file)
    bag_path = config.get("last_bag", "")
    if not bag_path:
        print("No bag file found in the configuration file.")
        exit(1)
    folder_path = os.path.dirname(bag_path)
    result_file = os.path.join(folder_path, "result.pdf")
    main(folder_path, result_file)
